#include "codegen.hpp"
#include "ast.hpp"
#include "opcode.hpp"
#include "symtable.hpp"
#include "codeobj.hpp"

#include <cassert>
#include <iostream>

using namespace std;

void Label::resolve(vector<unsigned char>& codes) {
    _pos = (unsigned short)codes.size();
    for (auto it = _uses.begin(); it != _uses.end(); it++) {
        int old_value = (((unsigned int)codes[*it-1]) << 8) 
            | ((unsigned int)codes[*it-2]);

        int param = 0;
        if (old_value == 1) {// absolute jump
            param = _pos;
        }
        else {
            param = _pos - *it;
        }
        codes[*it-1] = (unsigned char)((param >> 8) & 0xff);
        codes[*it-2] = (unsigned char)(param & 0xff);
    }

    _resolved = true;
}

CodeGen::CodeGen(SymbolTable* symtable, bool is_func) : _offset(0), _stack_size(0), _nlocals(0), _arg_count(0),
    _flags(0), _sym_table(symtable), _is_func(is_func) {
    _context = Context::Load;
}

CodeGen::~CodeGen() {
}

void CodeGen::prelogue(Node* root) {
    _type_checker->visit(root);
    if (_type_checker->has_error()) {
        exit(-1);
    }
}

void CodeGen::visit(Node* n) {
    n->accept(this);
}

void CodeGen::add_op(unsigned char op_code, unsigned short param) {
    _insts.push_back(op_code);
    _offset += 1;
    if (op_code >= HAVE_ARGUMENT) {
        _insts.push_back((unsigned char)(param & 0xff));
        _insts.push_back((unsigned char)((param >> 8) & 0xff));
        _offset += 2;
    }
}

void CodeGen::jump_abs(unsigned char op_code, Label& label) {
    if (label.is_resolved()) {
        add_op(op_code, label.pos());
    }
    else {
        add_op(op_code, 1); // 1 means absolute jump
        label.add_use(_offset);
    }
}

void CodeGen::jump_rel(unsigned char op_code, Label& label) {
    if (label.is_resolved()) {
        add_op(op_code, label.pos() - _offset);
    }
    else {
        add_op(op_code, 0); // 1 means absolute jump
        label.add_use(_offset);
    }
}

void CodeGen::bind(Label& label) {
    label.resolve(_insts);
}

void CodeGen::visit(PrintNode* n) {
    visit(n->value());
    add_op(PRINT_ITEM);
    add_op(PRINT_NEWLINE);
}

void CodeGen::visit(ReturnNode* n) {
    visit(n->value());
    add_op(RETURN_VALUE);
}

void CodeGen::visit(BinaryOp* op) {
    visit(op->left());
    visit(op->right());

    switch(op->op_type()) {
    case AST_OP_ADD:
        add_op(BINARY_ADD);
        break;
    case AST_OP_SUB:
        add_op(BINARY_SUB);
        break;
    case AST_OP_MUL:
        add_op(BINARY_MUL);
        break;
    case AST_OP_DIV:
        add_op(BINARY_DIV);
        break;
    case AST_OP_BIT_AND:
        add_op(BINARY_AND);
        break;
    case AST_OP_BIT_OR:
        add_op(BINARY_OR);
        break;
    case AST_OP_BIT_XOR:
        add_op(BINARY_XOR);
        break;
    case AST_OP_LOG_OR:
	    printf("(or ");
        break;
    case AST_OP_LOG_AND:
	    printf("(and ");
        break;
    case AST_OP_LEFT_SHIFT:
        add_op(BINARY_LSHIFT);
        break;
    case AST_OP_RIGHT_SHIFT:
        add_op(BINARY_RSHIFT);
        break;
    default:
        printf("Unknown binary op %d\n", op->op_type());
        return;
    }
}

void CodeGen::visit(ListNode* n) {
    for (auto it = n->node_list()->begin(); 
        it != n->node_list()->end(); it++) {
        visit(*it);
    }
}

unsigned short CodeGen::put_consts_if_absent(Constant& c) {
    auto it = _const_map.find(c);

    int index = -1;
    if (it == _const_map.end()) {
        index = _const_map.size();
        _const_map[c] = index;
        _consts.push_back(c);
    }
    else {
        index = it->second;
    }

    return (unsigned short)index;
}

void CodeGen::visit(ConstInt* n) {
    Constant c(VT_INT, n->_value);
    unsigned short index = put_consts_if_absent(c);
    add_op(LOAD_CONST, index);
}

void CodeGen::visit(ConstString* n) {
    Constant c(VT_STRING, n->_value);
    unsigned short index = put_consts_if_absent(c);
    add_op(LOAD_CONST, index);
}

void CodeGen::visit(ConstArray* n) {
    for (auto it = n->arr->begin(); it != n->arr->end(); it++) {
        visit(*it);
    }
    add_op(BUILD_LIST, n->arr->size());
}

void CodeGen::visit(BuiltinConst* n) {
    Constant* c = nullptr;
    unsigned short index = 0;
    switch (n->_builtin_id) {
    case BuiltinConst::UNIT:
        c = new Constant(VT_UNIT);
        index = put_consts_if_absent(*c);
        add_op(LOAD_CONST, index);
        break;
    case BuiltinConst::TRUE:
        c = new Constant(VT_TRUE);
        index = put_consts_if_absent(*c);
        add_op(LOAD_CONST, index);
        break;
    case BuiltinConst::FALSE:
        c = new Constant(VT_FALSE);
        index = put_consts_if_absent(*c);
        put_consts_if_absent(*c);
        break;
    default:
        cerr << "unknown const: " << n->_builtin_id;
    }
}

void CodeGen::visit(IfNode* n) {
    Label else_part, end;

    visit(n->cond());

    if (n->else_part() == NULL) {
        jump_abs(POP_JUMP_IF_FALSE, end);
    }
    else {
        jump_abs(POP_JUMP_IF_FALSE, else_part);
    }

    visit(n->then_part());

    if (n->else_part() != NULL) {
        jump_rel(JUMP_FORWARD, end);
        bind(else_part);
        visit(n->else_part());
    }
    
    bind(end);
}

void CodeGen::visit(WhileNode* n) {
    Label loop_begin, loop_end, block_end;

    jump_rel(SETUP_LOOP, block_end);
    bind(loop_begin);
    visit(n->cond());
    jump_abs(POP_JUMP_IF_FALSE, loop_end);
    visit(n->body());
    jump_abs(JUMP_ABSOLUTE, loop_begin);
    bind(loop_end);
    add_op(POP_BLOCK);
    bind(block_end);
}

int CodeGen::compare_op(int op) {
    switch (op) {
    case CmpNode::CMP_LT:
        return LESS;
    case CmpNode::CMP_EQU:
        return EQUAL;
    case CmpNode::CMP_GT:
        return GREATER;
    case CmpNode::CMP_GE:
        return GREATER_EQUAL;
    case CmpNode::CMP_LE:
        return LESS_EQUAL;
    case CmpNode::CMP_NE:
        return NOT_EQUAL;
    }

    return 0;
}

void CodeGen::visit(CmpNode* n) {
    visit(n->left());
    visit(n->right());
    add_op(COMPARE_OP, compare_op(n->cmp_op()));
}

void CodeGen::visit(LogicNotNode* n) {
    visit(n->value());
    add_op(UNARY_NOT);
}

void CodeGen::visit(AssignNode* n) {
    visit(n->right());
    _context = Context::Store;
    visit(n->left());
    _context = Context::Load;
}

void CodeGen::visit(VarNode* n) {
    int index = _sym_table->get_name_index(n->name());
    if (index >= 0) {
        if (_context == Context::Load) {
            add_op(LOAD_NAME, (unsigned short)index);
        }
        else {
            add_op(STORE_NAME, (unsigned short)index);
        }
        return;
    }

    index = _sym_table->get_free_var_index(n->name());
    if (index >= 0) {
        index += _sym_table->cellvars().size();
        if (_context == Context::Load) {
            add_op(LOAD_DEREF, (unsigned short)index);
        }
        else {
            add_op(STORE_DEREF, (unsigned short)index);
        }
        return;
    }

    index = _sym_table->get_cell_var_index(n->name());
    if (index >= 0) {
        if (_context == Context::Load) {
            add_op(LOAD_DEREF, (unsigned short)index);
        }
        else {
            add_op(STORE_DEREF, (unsigned short)index);
        }
        return;
    }

    index = _sym_table->get_var_name_index(n->name());
    if (index >= 0) {
        if (_context == Context::Load) {
            add_op(LOAD_FAST, (unsigned short)index);
        }
        else {
            add_op(STORE_FAST, (unsigned short)index);
        }
        return;
    }
}

void CodeGen::visit(ParamDefNode* n) {
}

void CodeGen::visit(VarDefNode* n) {
    if (n->init_value()) {
        visit(n->init_value());
    }

    int index = _sym_table->get_name_index(n->name());
    if (index >= 0) {
        add_op(STORE_NAME, (unsigned short)index);
        return;
    }

    index = _sym_table->get_var_name_index(n->name());
    if (index >= 0) {
        add_op(STORE_FAST, (unsigned short)index);
    }

    index = _sym_table->get_cell_var_index(n->name());
    if (index >= 0) {
        add_op(STORE_DEREF, (unsigned short)index);
    }

    // var def should not be free var.
}

void CodeGen::visit(CallNode* n) {
    visit(n->callee());
    visit(n->args());
    add_op(CALL_FUNCTION, (unsigned short)n->args()->size());
}

void CodeGen::visit(SubscrNode* n) {
    bool is_store = (_context == Context::Store);
    _context = Context::Load;

    // For generic instantiation.
    if (n->is_generic_inst()) {
        visit(n->name());
        return;
    }

    visit(n->name());
    visit(n->index());
    if (is_store) {
        add_op(STORE_SUBSCR);
    }
    else {
        add_op(BINARY_SUBSCR);
    }
}

void CodeGen::visit(FuncDefNode* n) {
    CodeGen func_gen(n->sym_table(), true);
    func_gen.set_arg_count(n->args()->size());
    func_gen.visit(n->body());

    Constant c(VT_CODE_OBJECT, func_gen.make_code_object());
    int index = _const_map.size();
    _const_map[c] = index;
    _consts.push_back(c);

    if (n->sym_table()->freevars().size() > 0) {
        for (auto it = n->sym_table()->freevars().begin();
            it != n->sym_table()->freevars().end(); it++) {
            int cnt = n->sym_table()->next->get_cell_var_index(*it);
            if (cnt >= 0) {
                add_op(LOAD_CLOSURE, (unsigned short)cnt);
            }
            else {
                cnt = n->sym_table()->next->cellvars().size();
                cnt += n->sym_table()->next->get_free_var_index(*it);
                add_op(LOAD_CLOSURE, (unsigned short)cnt);
            }
        }

        add_op(BUILD_TUPLE, (unsigned short)n->sym_table()->freevars().size());
        add_op(LOAD_CONST, (unsigned short)index);
        add_op(MAKE_CLOSURE, 0);
    }
    else {
        add_op(LOAD_CONST, (unsigned short)index);
        add_op(MAKE_FUNCTION, 0);
    }

    index = _sym_table->get_name_index(n->name());
    if (index >= 0) {
        add_op(STORE_NAME, (unsigned short)index);
        return;
    }

    index = _sym_table->get_var_name_index(n->name());
    if (index >= 0) {
        add_op(STORE_FAST, (unsigned short)index);
    }
}

void CodeGen::visit(LambdaDef* n) {
    CodeGen func_gen(n->sym_table(), true);
    func_gen.set_arg_count(1);
    func_gen.visit(n->body());

    Constant c(VT_CODE_OBJECT, func_gen.make_code_object());
    int index = _const_map.size();
    _const_map[c] = index;
    _consts.push_back(c);

    if (n->sym_table()->freevars().size() > 0) {
        for (auto it = n->sym_table()->freevars().begin();
            it != n->sym_table()->freevars().end(); it++) {
            int cnt = n->sym_table()->next->get_cell_var_index(*it);
            if (cnt >= 0) {
                add_op(LOAD_CLOSURE, (unsigned short)cnt);
            }
            else {
                cnt = n->sym_table()->next->cellvars().size();
                cnt += n->sym_table()->next->get_free_var_index(*it);
                add_op(LOAD_CLOSURE, (unsigned short)cnt);
            }
        }

        add_op(BUILD_TUPLE, (unsigned short)n->sym_table()->freevars().size());
        add_op(LOAD_CONST, (unsigned short)index);
        add_op(MAKE_CLOSURE, 0);
    }
    else {
        add_op(LOAD_CONST, (unsigned short)index);
        add_op(MAKE_FUNCTION, 0);
    }
}

void CodeGen::visit(TypeVarNode* n) {
    // do nothing.
}

void CodeGen::visit(ArrowTypeNode* n) {
    // do nothing.
}

void CodeGen::visit(TypeConstInt* n) {
    // do nothing
}

void CodeGen::visit(ArrayTypeNode* n) {
    // do nothing.
}

void CodeGen::visit(TypeVarDefNode* n) {
    // do nothing.
}

void CodeGen::visit(TypeFuncDefNode* n) {
    // do nothing.
}

void CodeGen::visit(TypeCallNode* n) {
    // do nothing.
}

CodeObject* CodeGen::make_code_object() {
    return new CodeObject(_arg_count, _nlocals, _stack_size, _flags, _insts, _consts, 
        _sym_table->names(), _sym_table->var_names(), _sym_table->freevars(),
        _sym_table->cellvars());
}

void Disassembler::print_detail(unsigned char c, unsigned short param) {
    if (c == COMPARE_OP) {
        switch(param) {
        case LESS:
            printf("(<)");
            break;
        case EQUAL:
            printf("(==)");
            break;
        case LESS_EQUAL:
            printf("(<=)");
            break;
        default:
            printf("(unknown compare operator : %d)", param);
        }
    }
}

void Disassembler::dis(vector<unsigned char>& insts) {
    int pc = 0;
    for (auto it = insts.begin(); it != insts.end(); it++) {
        unsigned char c = *it;
        if (c >= HAVE_ARGUMENT) {
            printf("%3d ", pc);
            printf("%s\t\t", bytecodes[c].c_str());
            unsigned short param = ((unsigned short)*(++it)) |
                    (((unsigned short)*(++it)) << 8);
            printf("%d", param);
            print_detail(c, param);
            printf("\n");
            pc += 3;
        }
        else {
            printf("%3d ", pc);
            printf("%s\n", bytecodes[c].c_str());
            pc += 1;
        }
    }
}

